package org.example.cases.sf1;

import com.antgroup.geaflow.api.function.base.FilterFunction;
import com.antgroup.geaflow.api.function.base.MapFunction;
import com.antgroup.geaflow.api.graph.PGraphWindow;
import com.antgroup.geaflow.api.graph.compute.VertexCentricCompute;
import com.antgroup.geaflow.api.graph.function.vc.VertexCentricCombineFunction;
import com.antgroup.geaflow.api.graph.function.vc.VertexCentricComputeFunction;
import com.antgroup.geaflow.api.pdata.stream.window.PWindowSource;
import com.antgroup.geaflow.api.pdata.stream.window.PWindowStream;
import com.antgroup.geaflow.api.window.impl.AllWindow;
import com.antgroup.geaflow.common.tuple.Tuple;
import com.antgroup.geaflow.model.graph.edge.EdgeDirection;
import com.antgroup.geaflow.model.graph.edge.IEdge;
import com.antgroup.geaflow.model.graph.edge.impl.ValueEdge;
import com.antgroup.geaflow.model.graph.vertex.IVertex;
import com.antgroup.geaflow.model.graph.vertex.impl.ValueVertex;
import com.antgroup.geaflow.pipeline.Pipeline;
import com.antgroup.geaflow.pipeline.task.PipelineTask;
import com.antgroup.geaflow.view.GraphViewBuilder;
import com.antgroup.geaflow.view.IViewDesc;
import com.antgroup.geaflow.view.graph.GraphViewDesc;
import org.example.constants.Constants;
import org.example.function.AbstractVcFunc;
import org.example.function.CSVFileSinkV2;
import org.example.function.CSVFileSourceV2;
import org.example.model.LongStringResult;
import org.example.util.FileUtils;

import java.util.Arrays;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;

public class Case4Pipeline {

    // private static final AndFilter edgeFilter =
    //         OutEdgeFilter.instance()
    //                 .and(new EdgeLabelFilter<>(Constants.EDGE_GUARANTEE));

    private final String snapshotPath;

    private final Pipeline pipeline;

    public Case4Pipeline(String snapshotPath, Pipeline pipeline) {
        this.snapshotPath = snapshotPath;
        this.pipeline = pipeline;
    }

    public void submit() {

        pipeline.submit((PipelineTask) pipelineTaskCxt -> {

            CSVFileSourceV2<IVertex<Long, Tuple<Double, Double>>> loanSource = new CSVFileSourceV2<>(
                    FileUtils.pathJoin(this.snapshotPath, Constants.FILE_LOAN),
                    Arrays.asList(Constants.LOAN_ID, Constants.LOAN_AMOUNT),
                    (CSVFileSourceV2.CSVLineParser<IVertex<Long, Tuple<Double, Double>>>) line -> {
                        ValueVertex<Long, Tuple<Double, Double>> vertex = new ValueVertex<>(
                                Long.valueOf(line[0]),
                                Tuple.of(Double.valueOf(line[1]), 0.0)
                        );
                        return Collections.singletonList(vertex);
                    }
            );

            PWindowSource<IVertex<Long, Tuple<Double, Double>>> loanVertices =
                    pipelineTaskCxt.buildSource(loanSource, AllWindow.getInstance())
                            .withParallelism(Constants.SOURCE_PARALLELISM);

            CSVFileSourceV2<IVertex<Long, Tuple<Double, Double>>> personSource = new CSVFileSourceV2<>(
                    FileUtils.pathJoin(this.snapshotPath, Constants.FILE_PERSON),
                    Collections.singletonList(Constants.PERSON_ID),
                    (CSVFileSourceV2.CSVLineParser<IVertex<Long, Tuple<Double, Double>>>) line -> {
                        ValueVertex<Long, Tuple<Double, Double>> vertex =
                                new ValueVertex<>(
                                        Long.valueOf(line[0]),
                                        Tuple.of(0.0, 0.0)
                                );
                        return Collections.singletonList(vertex);
                    }
            );

            PWindowSource<IVertex<Long, Tuple<Double, Double>>> personVertices =
                    pipelineTaskCxt.buildSource(personSource, AllWindow.getInstance())
                            .withParallelism(Constants.SOURCE_PARALLELISM);

            PWindowStream<IVertex<Long, Tuple<Double, Double>>> unionVerticesStream =
                    loanVertices.union(personVertices)
                            .withParallelism(Constants.UNION_PARALLELISM);

            CSVFileSourceV2<IEdge<Long, Integer>> personLoanSource = new CSVFileSourceV2<>(
                    FileUtils.pathJoin(this.snapshotPath, Constants.FILE_PERSON_APPLY_LOAN),
                    Arrays.asList(Constants.PERSON_ID, Constants.LOAN_ID),
                    (CSVFileSourceV2.CSVLineParser<IEdge<Long, Integer>>) line -> {
                        ValueEdge<Long, Integer> edge =
                                new ValueEdge<>(
                                        Long.valueOf(line[1]),
                                        Long.valueOf(line[0]),
                                        0,
                                        EdgeDirection.IN
                                );
                        return Collections.singletonList(edge);
                    }
            );
            PWindowSource<IEdge<Long, Integer>> personApplyLoanEdges = pipelineTaskCxt
                    .buildSource(personLoanSource, AllWindow.getInstance())
                    .withParallelism(Constants.SOURCE_PARALLELISM);

            CSVFileSourceV2<IEdge<Long, Integer>> personGuaranteeSource = new CSVFileSourceV2<>(
                    FileUtils.pathJoin(this.snapshotPath, Constants.FILE_PERSON_GUARANTEE_PERSON),
                    Arrays.asList(Constants.FROM_ID, Constants.TO_ID),
                    (CSVFileSourceV2.CSVLineParser<IEdge<Long, Integer>>) line -> {
                        ValueEdge<Long, Integer> edge =
                                new ValueEdge<>(
                                        Long.valueOf(line[0]),
                                        Long.valueOf(line[1]),
                                        1
                                );
                        return Collections.singletonList(edge);
                    }
            );
            PWindowSource<IEdge<Long, Integer>> personGuaranteePersonEdges = pipelineTaskCxt
                    .buildSource(personGuaranteeSource, AllWindow.getInstance())
                    .withParallelism(Constants.SOURCE_PARALLELISM);

            PWindowStream<IEdge<Long, Integer>> unionEdgesStream =
                    personApplyLoanEdges.union(personGuaranteePersonEdges)
                            .withParallelism(Constants.UNION_PARALLELISM);

            int iterationParallelism = Constants.ITERATOR_PARALLELISM;
            GraphViewDesc graphViewDesc = GraphViewBuilder
                    .createGraphView("case4")
                    .withShardNum(iterationParallelism)
                    .withBackend(IViewDesc.BackendType.Memory)
                    .build();
            PGraphWindow<Long, Tuple<Double, Double>, Integer> graphWindow =
                    pipelineTaskCxt.buildWindowStreamGraph(unionVerticesStream, unionEdgesStream, graphViewDesc);

            graphWindow.compute(new Case4VCAlgorithm(6))
                    .compute(iterationParallelism)
                    .getVertices()
                    .filter((FilterFunction<IVertex<Long, Tuple<Double, Double>>>) record ->
                            record.getValue().f0 >= 0 && record.getValue().f1 > 0)
                    .withParallelism(Constants.FILTER_PARALLELISM)
                    .map((MapFunction<IVertex<Long, Tuple<Double, Double>>, LongStringResult>) valueLabelVertex ->
                            new LongStringResult(
                                    valueLabelVertex.getId(),
                                    String.format("%.2f", valueLabelVertex.getValue().f1 / 100_000_000))
                    )
                    .withParallelism(Constants.MAP_PARALLELISM)
                    .sink(new CSVFileSinkV2<>("result4.csv"))
                    .withParallelism(Constants.SINK_PARALLELISM);
        });
    }

    public static class Case4VCAlgorithm extends VertexCentricCompute<Long, Tuple<Double, Double>, Integer, Double> {

        public Case4VCAlgorithm(long iterations) {
            super(iterations);
        }

        @Override
        public VertexCentricComputeFunction<Long, Tuple<Double, Double>, Integer, Double> getComputeFunction() {
            return new Case4VertexCentricComputeFunction();
        }

        @Override
        public VertexCentricCombineFunction<Double> getCombineFunction() {
            return null;
        }
    }

    public static class Case4VertexCentricComputeFunction extends AbstractVcFunc<Long, Tuple<Double, Double>, Integer, Double> {

        @Override
        public void compute(Long vertexId, Iterator<Double> messageIterator) {

            long iterationId = this.context.getIterationId();
            Tuple<Double, Double> vertexValue = this.context.vertex().get().getValue();

            if (iterationId == 1) {
                if (vertexValue.f0 > 0) {
                    this.context.sendMessageToNeighbors(vertexValue.f0);
                    vertexValue.f0 = -vertexValue.f0;
                } else {
                    if (!this.context.edges().getOutEdges().isEmpty()) {
                        this.context.sendMessage(vertexId, 0.0);
                    }
                }
            } else if (iterationId == 2) {
                double sum = 0.0;
                while (messageIterator.hasNext()) {
                    double msg = messageIterator.next();
                    sum += msg;
                }
                vertexValue.f0 = sum;
                this.context.setNewVertexValue(vertexValue);
                List<IEdge<Long, Integer>> personOutEdges = this.context.edges().getOutEdges();
                for (IEdge<Long, Integer> outEdge : personOutEdges) {
                    if (outEdge.getValue() == 1) {
                        this.context.sendMessage(outEdge.getTargetId(), (double) -vertexId);
                    }
                }
            } else if (iterationId >= 3 && iterationId <= 5) {
                List<IEdge<Long, Integer>> outEdges = this.context.edges().getOutEdges();
                while (messageIterator.hasNext()) {
                    double msg = messageIterator.next();
                    if (msg >= 0) {
                        vertexValue.f1 += msg;
                    } else {
                        this.context.sendMessage((long) (-msg), vertexValue.f0);
                    }
                    for (IEdge<Long, Integer> outEdge : outEdges) {
                        if (msg < 0 && outEdge.getValue() == 1) {
                            this.context.sendMessage(outEdge.getTargetId(), msg);
                        }
                    }
                }
                this.context.setNewVertexValue(vertexValue);
            } else {
                while (messageIterator.hasNext()) {
                    double msg = messageIterator.next();
                    if (msg > 0) {
                        vertexValue.f1 += msg;
                    }
                }
                this.context.setNewVertexValue(vertexValue);
            }
        }
    }
}
